import {
  EnhancedGenerateContentResponse,
  Content,
  Part,
  type FunctionDeclarationsTool as GoogleGenerativeAIFunctionDeclarationsTool,
  type FunctionDeclaration as GenerativeAIFunctionDeclaration,
  POSSIBLE_ROLES,
  FunctionCallPart,
  TextPart,
  FileDataPart,
  InlineDataPart,
  type GenerateContentResponse,
} from "@google/generative-ai";
import {
  AIMessage,
  AIMessageChunk,
  BaseMessage,
  ChatMessage,
  ToolMessage,
  ToolMessageChunk,
  MessageContent,
  MessageContentComplex,
  UsageMetadata,
  isAIMessage,
  isBaseMessage,
  isToolMessage,
  StandardContentBlockConverter,
  parseBase64DataUrl,
  convertToProviderContentBlock,
  isDataContentBlock,
  InputTokenDetails,
} from "@langchain/core/messages";
import {
  ChatGeneration,
  ChatGenerationChunk,
  ChatResult,
} from "@langchain/core/outputs";
import { isLangChainTool } from "@langchain/core/utils/function_calling";
import { isOpenAITool } from "@langchain/core/language_models/base";
import { ToolCallChunk } from "@langchain/core/messages/tool";
import { v4 as uuidv4 } from "uuid";
import {
  jsonSchemaToGeminiParameters,
  schemaToGenerativeAIParameters,
} from "./zod_to_genai_parameters.js";
import { GoogleGenerativeAIToolType } from "../types.js";

export const _FUNCTION_CALL_THOUGHT_SIGNATURES_MAP_KEY =
  "__gemini_function_call_thought_signatures__";
const DUMMY_SIGNATURE =
  "ErYCCrMCAdHtim9kOoOkrPiCNVsmlpMIKd7ZMxgiFbVQOkgp7nlLcDMzVsZwIzvuT7nQROivoXA72ccC2lSDvR0Gh7dkWaGuj7ctv6t7ZceHnecx0QYa+ix8tYpRfjhyWozQ49lWiws6+YGjCt10KRTyWsZ2h6O7iHTYJwKIRwGUHRKy/qK/6kFxJm5ML00gLq4D8s5Z6DBpp2ZlR+uF4G8jJgeWQgyHWVdx2wGYElaceVAc66tZdPQRdOHpWtgYSI1YdaXgVI8KHY3/EfNc2YqqMIulvkDBAnuMhkAjV9xmBa54Tq+ih3Im4+r3DzqhGqYdsSkhS0kZMwte4Hjs65dZzCw9lANxIqYi1DJ639WNPYihp/DCJCos7o+/EeSPJaio5sgWDyUnMGkY1atsJZ+m7pj7DD5tvQ==";

const iife = (fn: () => string) => fn();

export function getMessageAuthor(message: BaseMessage) {
  if (ChatMessage.isInstance(message)) {
    return message.role;
  }
  return message.type;
}

/**
 * Maps a message type to a Google Generative AI chat author.
 * @param message The message to map.
 * @param model The model to use for mapping.
 * @returns The message type mapped to a Google Generative AI chat author.
 */
export function convertAuthorToRole(
  author: string
): (typeof POSSIBLE_ROLES)[number] {
  switch (author) {
    /**
     *  Note: Gemini currently is not supporting system messages
     *  we will convert them to human messages and merge with following
     * */
    case "supervisor":
    case "ai":
    case "model":
      return "model";
    case "system":
      return "system";
    case "human":
      return "user";
    case "tool":
    case "function":
      return "function";
    default:
      throw new Error(`Unknown / unsupported author: ${author}`);
  }
}

function messageContentMedia(content: MessageContentComplex): Part {
  if ("mimeType" in content && "data" in content) {
    return {
      inlineData: {
        mimeType: content.mimeType,
        data: content.data,
      },
    };
  }
  if ("mimeType" in content && "fileUri" in content) {
    return {
      fileData: {
        mimeType: content.mimeType,
        fileUri: content.fileUri,
      },
    };
  }

  throw new Error("Invalid media content");
}

function inferToolNameFromPreviousMessages(
  message: ToolMessage | ToolMessageChunk,
  previousMessages: BaseMessage[]
): string | undefined {
  return previousMessages
    .map((msg) => {
      if (isAIMessage(msg)) {
        return msg.tool_calls ?? [];
      }
      return [];
    })
    .flat()
    .find((toolCall) => {
      return toolCall.id === message.tool_call_id;
    })?.name;
}

function _getStandardContentBlockConverter(isMultimodalModel: boolean) {
  const standardContentBlockConverter: StandardContentBlockConverter<{
    text: TextPart;
    image: FileDataPart | InlineDataPart;
    audio: FileDataPart | InlineDataPart;
    file: FileDataPart | InlineDataPart | TextPart;
  }> = {
    providerName: "Google Gemini",

    fromStandardTextBlock(block) {
      return {
        text: block.text,
      };
    },

    fromStandardImageBlock(block): FileDataPart | InlineDataPart {
      if (!isMultimodalModel) {
        throw new Error("This model does not support images");
      }
      if (block.source_type === "url") {
        const data = parseBase64DataUrl({ dataUrl: block.url });
        if (data) {
          return {
            inlineData: {
              mimeType: data.mime_type,
              data: data.data,
            },
          };
        } else {
          return {
            fileData: {
              mimeType: block.mime_type ?? "",
              fileUri: block.url,
            },
          };
        }
      }

      if (block.source_type === "base64") {
        return {
          inlineData: {
            mimeType: block.mime_type ?? "",
            data: block.data,
          },
        };
      }

      throw new Error(`Unsupported source type: ${block.source_type}`);
    },

    fromStandardAudioBlock(block): FileDataPart | InlineDataPart {
      if (!isMultimodalModel) {
        throw new Error("This model does not support audio");
      }
      if (block.source_type === "url") {
        const data = parseBase64DataUrl({ dataUrl: block.url });
        if (data) {
          return {
            inlineData: {
              mimeType: data.mime_type,
              data: data.data,
            },
          };
        } else {
          return {
            fileData: {
              mimeType: block.mime_type ?? "",
              fileUri: block.url,
            },
          };
        }
      }

      if (block.source_type === "base64") {
        return {
          inlineData: {
            mimeType: block.mime_type ?? "",
            data: block.data,
          },
        };
      }

      throw new Error(`Unsupported source type: ${block.source_type}`);
    },

    fromStandardFileBlock(block): FileDataPart | InlineDataPart | TextPart {
      if (!isMultimodalModel) {
        throw new Error("This model does not support files");
      }
      if (block.source_type === "text") {
        return {
          text: block.text,
        };
      }
      if (block.source_type === "url") {
        const data = parseBase64DataUrl({ dataUrl: block.url });
        if (data) {
          return {
            inlineData: {
              mimeType: data.mime_type,
              data: data.data,
            },
          };
        } else {
          return {
            fileData: {
              mimeType: block.mime_type ?? "",
              fileUri: block.url,
            },
          };
        }
      }

      if (block.source_type === "base64") {
        return {
          inlineData: {
            mimeType: block.mime_type ?? "",
            data: block.data,
          },
        };
      }
      throw new Error(`Unsupported source type: ${block.source_type}`);
    },
  };
  return standardContentBlockConverter;
}

function _convertLangChainContentToPart(
  content: MessageContentComplex,
  isMultimodalModel: boolean
): Part | undefined {
  if (isDataContentBlock(content)) {
    return convertToProviderContentBlock(
      content,
      _getStandardContentBlockConverter(isMultimodalModel)
    );
  }

  if (content.type === "text") {
    return { text: content.text };
  } else if (content.type === "executableCode") {
    return { executableCode: content.executableCode };
  } else if (content.type === "codeExecutionResult") {
    return { codeExecutionResult: content.codeExecutionResult };
  } else if (content.type === "image_url") {
    if (!isMultimodalModel) {
      throw new Error(`This model does not support images`);
    }
    let source;
    if (typeof content.image_url === "string") {
      source = content.image_url;
    } else if (
      typeof content.image_url === "object" &&
      "url" in content.image_url
    ) {
      source = content.image_url.url;
    } else {
      throw new Error("Please provide image as base64 encoded data URL");
    }
    const [dm, data] = source.split(",");
    if (!dm.startsWith("data:")) {
      throw new Error("Please provide image as base64 encoded data URL");
    }

    const [mimeType, encoding] = dm.replace(/^data:/, "").split(";");
    if (encoding !== "base64") {
      throw new Error("Please provide image as base64 encoded data URL");
    }

    return {
      inlineData: {
        data,
        mimeType,
      },
    };
  } else if (content.type === "media") {
    return messageContentMedia(content);
  } else if (content.type === "tool_use") {
    return {
      functionCall: {
        name: content.name,
        args: content.input,
      },
    };
  } else if (
    content.type?.includes("/") &&
    // Ensure it's a single slash.
    content.type.split("/").length === 2 &&
    "data" in content &&
    typeof content.data === "string"
  ) {
    return {
      inlineData: {
        mimeType: content.type,
        data: content.data,
      },
    };
  } else if ("functionCall" in content) {
    // No action needed here — function calls will be added later from message.tool_calls
    return undefined;
  } else {
    if ("type" in content) {
      throw new Error(`Unknown content type ${content.type}`);
    } else {
      throw new Error(`Unknown content ${JSON.stringify(content)}`);
    }
  }
}

export function convertMessageContentToParts(
  message: BaseMessage,
  isMultimodalModel: boolean,
  previousMessages: BaseMessage[],
  model?: string
): Part[] {
  if (isToolMessage(message)) {
    const messageName =
      message.name ??
      inferToolNameFromPreviousMessages(message, previousMessages);
    if (messageName === undefined) {
      throw new Error(
        `Google requires a tool name for each tool call response, and we could not infer a called tool name for ToolMessage "${message.id}" from your passed messages. Please populate a "name" field on that ToolMessage explicitly.`
      );
    }

    const result = Array.isArray(message.content)
      ? (message.content
          .map((c) => _convertLangChainContentToPart(c, isMultimodalModel))
          .filter((p) => p !== undefined) as Part[])
      : message.content;

    if (message.status === "error") {
      return [
        {
          functionResponse: {
            name: messageName,
            // The API expects an object with an `error` field if the function call fails.
            // `error` must be a valid object (not a string or array), so we wrap `message.content` here
            response: { error: { details: result } },
          },
        },
      ];
    }

    return [
      {
        functionResponse: {
          name: messageName,
          // again, can't have a string or array value for `response`, so we wrap it as an object here
          response: { result },
        },
      },
    ];
  }

  let functionCalls: FunctionCallPart[] = [];
  const messageParts: Part[] = [];

  if (typeof message.content === "string" && message.content) {
    messageParts.push({ text: message.content });
  }

  if (Array.isArray(message.content)) {
    messageParts.push(
      ...(message.content
        .map((c) => _convertLangChainContentToPart(c, isMultimodalModel))
        .filter((p) => p !== undefined) as Part[])
    );
  }

  const functionThoughtSignatures = message.additional_kwargs?.[
    _FUNCTION_CALL_THOUGHT_SIGNATURES_MAP_KEY
  ] as Record<string, string>;

  if (isAIMessage(message) && message.tool_calls?.length) {
    functionCalls = message.tool_calls.map((tc) => {
      const thoughtSignature = iife(() => {
        if (tc.id) {
          const signature = functionThoughtSignatures?.[tc.id];
          if (signature) {
            return signature;
          }
        }
        if (model?.includes("gemini-3")) {
          return DUMMY_SIGNATURE;
        }
        return "";
      });
      return {
        functionCall: {
          name: tc.name,
          args: tc.args,
        },
        ...(thoughtSignature ? { thoughtSignature } : {}),
      };
    });
  }

  return [...messageParts, ...functionCalls];
}

export function convertBaseMessagesToContent(
  messages: BaseMessage[],
  isMultimodalModel: boolean,
  convertSystemMessageToHumanContent: boolean = false,
  model: string
) {
  return messages.reduce<{
    content: Content[];
    mergeWithPreviousContent: boolean;
  }>(
    (acc, message, index) => {
      if (!isBaseMessage(message)) {
        throw new Error("Unsupported message input");
      }
      const author = getMessageAuthor(message);
      if (author === "system" && index !== 0) {
        throw new Error("System message should be the first one");
      }
      const role = convertAuthorToRole(author);

      const prevContent = acc.content[acc.content.length];
      if (
        !acc.mergeWithPreviousContent &&
        prevContent &&
        prevContent.role === role
      ) {
        throw new Error(
          "Google Generative AI requires alternate messages between authors"
        );
      }

      const parts = convertMessageContentToParts(
        message,
        isMultimodalModel,
        messages.slice(0, index),
        model
      );

      if (acc.mergeWithPreviousContent) {
        const prevContent = acc.content[acc.content.length - 1];
        if (!prevContent) {
          throw new Error(
            "There was a problem parsing your system message. Please try a prompt without one."
          );
        }
        prevContent.parts.push(...parts);

        return {
          mergeWithPreviousContent: false,
          content: acc.content,
        };
      }
      let actualRole = role;
      if (
        actualRole === "function" ||
        (actualRole === "system" && !convertSystemMessageToHumanContent)
      ) {
        // GenerativeAI API will throw an error if the role is not "user" or "model."
        actualRole = "user";
      }
      const content: Content = {
        role: actualRole,
        parts,
      };
      return {
        mergeWithPreviousContent:
          author === "system" && !convertSystemMessageToHumanContent,
        content: [...acc.content, content],
      };
    },
    { content: [], mergeWithPreviousContent: false }
  ).content;
}

export function mapGenerateContentResultToChatResult(
  response: EnhancedGenerateContentResponse,
  extra?: {
    usageMetadata: UsageMetadata | undefined;
  }
): ChatResult {
  // if rejected or error, return empty generations with reason in filters
  if (
    !response.candidates ||
    response.candidates.length === 0 ||
    !response.candidates[0]
  ) {
    return {
      generations: [],
      llmOutput: {
        filters: response.promptFeedback,
      },
    };
  }
  const [candidate] = response.candidates;
  const { content: candidateContent, ...generationInfo } = candidate;
  const functionCalls = candidateContent.parts?.reduce((acc, p) => {
    if ("functionCall" in p && p.functionCall) {
      acc.push({
        ...p,
        id:
          "id" in p.functionCall && typeof p.functionCall.id === "string"
            ? p.functionCall.id
            : uuidv4(),
      });
    }
    return acc;
  }, [] as (FunctionCallPart & { id: string })[]);
  let content: MessageContent | undefined;

  if (
    Array.isArray(candidateContent?.parts) &&
    candidateContent.parts.length === 1 &&
    candidateContent.parts[0].text
  ) {
    content = candidateContent.parts[0].text;
  } else if (
    Array.isArray(candidateContent?.parts) &&
    candidateContent.parts.length > 0
  ) {
    content = candidateContent.parts.map((p) => {
      if ("text" in p) {
        return {
          type: "text",
          text: p.text,
        };
      } else if ("inlineData" in p) {
        return {
          type: "inlineData",
          inlineData: p.inlineData,
        };
      } else if ("functionCall" in p) {
        return {
          type: "functionCall",
          functionCall: p.functionCall,
        };
      } else if ("functionResponse" in p) {
        return {
          type: "functionResponse",
          functionResponse: p.functionResponse,
        };
      } else if ("fileData" in p) {
        return {
          type: "fileData",
          fileData: p.fileData,
        };
      } else if ("executableCode" in p) {
        return {
          type: "executableCode",
          executableCode: p.executableCode,
        };
      } else if ("codeExecutionResult" in p) {
        return {
          type: "codeExecutionResult",
          codeExecutionResult: p.codeExecutionResult,
        };
      }
      return p;
    });
  } else {
    // no content returned - likely due to abnormal stop reason, e.g. malformed function call
    content = [];
  }

  const functionThoughtSignatures = functionCalls?.reduce((acc, fc) => {
    if ("thoughtSignature" in fc && typeof fc.thoughtSignature === "string") {
      acc[fc.id] = fc.thoughtSignature;
    }
    return acc;
  }, {} as Record<string, string>);

  let text = "";
  if (typeof content === "string") {
    text = content;
  } else if (Array.isArray(content) && content.length > 0) {
    const block = content.find((b) => "text" in b) as
      | { text: string }
      | undefined;
    text = block?.text ?? text;
  }

  const generation: ChatGeneration = {
    text,
    message: new AIMessage({
      content: content ?? "",
      tool_calls: functionCalls?.map((fc) => ({
        type: "tool_call",
        id: fc.id,
        name: fc.functionCall.name,
        args: fc.functionCall.args,
      })),
      additional_kwargs: {
        ...generationInfo,
        [_FUNCTION_CALL_THOUGHT_SIGNATURES_MAP_KEY]: functionThoughtSignatures,
      },
      usage_metadata: extra?.usageMetadata,
    }),
    generationInfo,
  };

  return {
    generations: [generation],
    llmOutput: {
      tokenUsage: {
        promptTokens: extra?.usageMetadata?.input_tokens,
        completionTokens: extra?.usageMetadata?.output_tokens,
        totalTokens: extra?.usageMetadata?.total_tokens,
      },
    },
  };
}

export function convertResponseContentToChatGenerationChunk(
  response: EnhancedGenerateContentResponse,
  extra: {
    usageMetadata?: UsageMetadata | undefined;
    index: number;
  }
): ChatGenerationChunk | null {
  if (!response.candidates || response.candidates.length === 0) {
    return null;
  }
  const [candidate] = response.candidates;
  const { content: candidateContent, ...generationInfo } = candidate;
  const functionCalls = candidateContent.parts?.reduce((acc, p) => {
    if ("functionCall" in p && p.functionCall) {
      acc.push({
        ...p,
        id:
          "id" in p.functionCall && typeof p.functionCall.id === "string"
            ? p.functionCall.id
            : uuidv4(),
      });
    }
    return acc;
  }, [] as (FunctionCallPart & { id: string })[]);
  let content: MessageContent | undefined;
  // Checks if some parts do not have text. If false, it means that the content is a string.
  if (
    Array.isArray(candidateContent?.parts) &&
    candidateContent.parts.every((p) => "text" in p)
  ) {
    content = candidateContent.parts.map((p) => p.text).join("");
  } else if (Array.isArray(candidateContent?.parts)) {
    content = candidateContent.parts.map((p) => {
      if ("text" in p) {
        return {
          type: "text",
          text: p.text,
        };
      } else if ("inlineData" in p) {
        return {
          type: "inlineData",
          inlineData: p.inlineData,
        };
      } else if ("functionCall" in p) {
        return {
          type: "functionCall",
          functionCall: p.functionCall,
        };
      } else if ("functionResponse" in p) {
        return {
          type: "functionResponse",
          functionResponse: p.functionResponse,
        };
      } else if ("fileData" in p) {
        return {
          type: "fileData",
          fileData: p.fileData,
        };
      } else if ("executableCode" in p) {
        return {
          type: "executableCode",
          executableCode: p.executableCode,
        };
      } else if ("codeExecutionResult" in p) {
        return {
          type: "codeExecutionResult",
          codeExecutionResult: p.codeExecutionResult,
        };
      }
      return p;
    });
  } else {
    // no content returned - likely due to abnormal stop reason, e.g. malformed function call
    content = [];
  }

  let text = "";
  if (content && typeof content === "string") {
    text = content;
  } else if (Array.isArray(content)) {
    const block = content.find((b) => "text" in b) as
      | { text: string }
      | undefined;
    text = block?.text ?? "";
  }

  const toolCallChunks: ToolCallChunk[] = [];
  if (functionCalls) {
    toolCallChunks.push(
      ...functionCalls.map((fc) => ({
        type: "tool_call_chunk" as const,
        id: fc.id,
        name: fc.functionCall.name,
        args: JSON.stringify(fc.functionCall.args),
      }))
    );
  }

  const functionThoughtSignatures = functionCalls?.reduce((acc, fc) => {
    if ("thoughtSignature" in fc && typeof fc.thoughtSignature === "string") {
      acc[fc.id] = fc.thoughtSignature;
    }
    return acc;
  }, {} as Record<string, string>);

  return new ChatGenerationChunk({
    text,
    message: new AIMessageChunk({
      content: content || "",
      name: !candidateContent ? undefined : candidateContent.role,
      tool_call_chunks: toolCallChunks,
      // Each chunk can have unique "generationInfo", and merging strategy is unclear,
      // so leave blank for now.
      additional_kwargs: {
        [_FUNCTION_CALL_THOUGHT_SIGNATURES_MAP_KEY]: functionThoughtSignatures,
      },
      response_metadata: {
        model_provider: "google-genai",
      },
      usage_metadata: extra.usageMetadata,
    }),
    generationInfo,
  });
}

export function convertToGenerativeAITools(
  tools: GoogleGenerativeAIToolType[]
): GoogleGenerativeAIFunctionDeclarationsTool[] {
  if (
    tools.every(
      (tool) =>
        "functionDeclarations" in tool &&
        Array.isArray(tool.functionDeclarations)
    )
  ) {
    return tools as GoogleGenerativeAIFunctionDeclarationsTool[];
  }
  return [
    {
      functionDeclarations: tools.map(
        (tool): GenerativeAIFunctionDeclaration => {
          if (isLangChainTool(tool)) {
            const jsonSchema = schemaToGenerativeAIParameters(tool.schema);
            if (
              jsonSchema.type === "object" &&
              "properties" in jsonSchema &&
              Object.keys(jsonSchema.properties).length === 0
            ) {
              return {
                name: tool.name,
                description: tool.description,
              };
            }
            return {
              name: tool.name,
              description: tool.description,
              parameters: jsonSchema,
            };
          }
          if (isOpenAITool(tool)) {
            return {
              name: tool.function.name,
              description:
                tool.function.description ?? `A function available to call.`,
              parameters: jsonSchemaToGeminiParameters(
                tool.function.parameters
              ),
            };
          }
          return tool as unknown as GenerativeAIFunctionDeclaration;
        }
      ),
    },
  ];
}

export function convertUsageMetadata(
  usageMetadata: GenerateContentResponse["usageMetadata"],
  model: string
): UsageMetadata {
  const output: UsageMetadata = {
    input_tokens: usageMetadata?.promptTokenCount ?? 0,
    output_tokens: usageMetadata?.candidatesTokenCount ?? 0,
    total_tokens: usageMetadata?.totalTokenCount ?? 0,
  };
  if (usageMetadata?.cachedContentTokenCount) {
    output.input_token_details ??= {};
    output.input_token_details.cache_read =
      usageMetadata.cachedContentTokenCount;
  }
  // gemini-3-pro-preview has bracket based tracking of tokens per request
  // FIXME(hntrl): move this usageMetadata calculation elsewhere
  if (model === "gemini-3-pro-preview") {
    const over200k = Math.max(0, usageMetadata?.promptTokenCount ?? 0 - 200000);
    const cachedOver200k = Math.max(
      0,
      usageMetadata?.cachedContentTokenCount ?? 0 - 200000
    );
    if (over200k) {
      output.input_token_details = {
        ...output.input_token_details,
        over_200k: over200k,
      } as InputTokenDetails;
    }
    if (cachedOver200k) {
      output.input_token_details = {
        ...output.input_token_details,
        cache_read_over_200k: cachedOver200k,
      } as InputTokenDetails;
    }
  }
  return output;
}
